import numpy as np
import numba as nb
from numba import cuda
from math import exp, ceil
from ..utils import get_grid_2d, assume_ndarray


@cuda.jit()
def _gpu_multiply(a: np.ndarray, b: np.ndarray, c: np.ndarray):
    x, y = cuda.grid(2)
    if x < a.shape[0] and y < a.shape[1]:
        c[x, y] = a[x, y] * b[y]


def multiply(a: np.ndarray, b: np.ndarray):
    a, b = assume_ndarray(a, b)
    if a.ndim < b.ndim:
        a, b = b, a

    assert a.shape[-b.ndim:] == b.shape, "%s与%s不匹配" % (a.shape, b.shape)

    shape = a.shape
    a = a.reshape((a.size // b.size, b.size))
    b = b.reshape((b.size,))
    c = np.zeros(a.shape)

    grid, block = get_grid_2d(a.shape[0], a.shape[1])
    _gpu_multiply[grid, block](a, b, c)

    return c.reshape(shape)
